import json
import math

Labels = {
    "Legend": 0,
    "CategoryTitle": 1,
    "Title": 2,
    "ValueTitle": 3,
    "BarCanvas": 4,
    "LineCanvas": 5,
    "PieCanvas": 6,
    "Bar": 7,
    "Line": 8,
    "Pie": 9,
    "Pictogram": 10
}


class PieBenchmark:
    def __init__(self, gt_label_path, pred_label_path=None):
        self.total_score = None
        self.rec = None
        self.total = None
        self.scores = None
        self.gt_path = gt_label_path
        self.pred_path = pred_label_path
        self.gt = dict()
        self.pred = dict()
        self.load()

    def load(self):
        with open(self.gt_path, "r") as fp:
            gt = json.load(fp)

        # 赛选gt中饼条形图数据元素
        for key in gt.keys():
            labels = gt[key]
            if "Pie" in labels.keys():
                self.gt[key] = labels["Pie"]

        with open(self.pred_path, "r") as fp:
            pred_data = json.load(fp)

        for absolute_path in pred_data:
            v = pred_data[absolute_path]
            filename = absolute_path.split("\\")[-1]
            self.pred[filename] = v

    @staticmethod
    def check(gt_bbox, pred_bbox, threshold):
        # 主要是比较两者中心是否差距很大
        gt_w = gt_bbox[2] - gt_bbox[0]
        gt_h = gt_bbox[3] - gt_bbox[1]

        diff_x = abs(pred_bbox[4]- gt_bbox[4])
        if diff_x > gt_w * threshold:
            return False

        diff_y = abs(pred_bbox[5]- gt_bbox[5])
        if diff_y > gt_h * threshold:
            return False

        return True

    def run(self, threshold=0.1):
        # 取每个gt与pred做对比
        scores = {}
        total_score = 0
        rec = 0
        total = 0
        for key in self.pred.keys():
            gt = self.gt[key]
            pred = self.pred[key]["data"]

            score = []
            for item in gt:
                total += 1
                pie_bbox_l = min(item[::2])
                pie_bbox_r = max(item[::2])
                pie_bbox_t = min(item[1::2])
                pie_bbox_b = max(item[1::2])
                pie_bbox_gt = [pie_bbox_l, pie_bbox_t,
                               pie_bbox_r, pie_bbox_b,
                               (pie_bbox_l + pie_bbox_r) / 2,
                               (pie_bbox_t + pie_bbox_b) / 2]

                cx = item[0]
                cy = item[1]
                x1 = item[2] - cx
                y1 = item[3] - cy
                x2 = item[4] - cx
                y2 = item[5] - cy

                theta_acos = (x1 * x2 + y1 * y2) / (math.sqrt(x1*x1 + y1*y1) * math.sqrt(x2*x2 + y2*y2))
                theta_gt = math.acos(theta_acos) / math.pi * 180

                max_score = -1

                for item_p in pred:
                    # [[theta_y, theta, [center_x, center_y], R, string_center, color, [x1, y1, x2, y2]]]
                    li = item_p[2][:]

                    p1p2 = item_p[-1]
                    li.append(li[0] + p1p2[0])
                    li.append(li[1] + p1p2[1])
                    li.append(li[0] + p1p2[2])
                    li.append(li[1] + p1p2[3])

                    pie_bbox_l_p = min(li[::2])
                    pie_bbox_r_p = max(li[::2])
                    pie_bbox_t_p = min(li[1::2])
                    pie_bbox_b_p = max(li[1::2])
                    pie_bbox_pred = [pie_bbox_l_p, pie_bbox_t_p,
                                     pie_bbox_r_p, pie_bbox_b_p,
                                     (pie_bbox_l_p + pie_bbox_r_p) / 2,
                                     (pie_bbox_t_p + pie_bbox_b_p) / 2]
                    theta_pred = item_p[1]

                    if self.check(pie_bbox_gt, pie_bbox_pred, threshold):
                        if theta_pred > 180:
                            tmp_score = 1 - abs((360 - theta_gt) - theta_pred) / (360 - theta_gt)
                        else:
                            tmp_score = 1 - abs(theta_gt - theta_pred) / theta_gt
                        if tmp_score > max_score:
                            max_score = tmp_score

                if max_score < 0:
                    score.append(0)
                    continue

                rec += 1
                score.append(max_score)
                total_score += max_score

            scores[key] = score
        self.total_score = total_score
        self.scores = scores
        self.total = total
        self.rec = rec

    def summary(self):
        if self.scores is None:
            return
        print(self.total, self.total_score, self.rec)
        print(f"识别率：{round(self.rec/self.total * 100, 4)}")
        print(f"召回率：{round(self.total_score/self.total * 100, 4)}")


if __name__ == "__main__":
    gt_labels = r"D:\Code\script\_banchmark\benchmark\dataset\CV\pie\labels\val.json"
    pred_labels = r"D:\Code\script\_banchmark\benchmark\utils\cornernet_pie_pred.json"
    benchmark = PieBenchmark(gt_labels, pred_labels)
    benchmark.run()
    benchmark.summary()
    # [[theta_y, theta, [center_x, center_y], R, string_center, color, [x1, y1, x2, y2]]]
